package org.com.sec.critical;

import java.math.BigInteger;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.RSAPublicKeySpec;
import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;

/**
 * Provides easy-to-use RSA asymmetric cryptography system.
 *
 * @author See <a href="the Google Code page of the easy-utils
 * library">http://code.google.com/p/easy-utils/</a> for collaborators and other
 * information.
 */
public abstract class RSA {

    private static BigInteger PUBLIC_MODULUS, PUBLIC_EXPONENT, PRIVATE_MODULUS, PRIVATE_EXPONENT;
    private static final int DEFAULT_KEY_SIZE = 2048;

    /**
     * Provides the modulus of the public key.
     *
     * @return {@link BigInteger} The modulus of the public key.
     */
    public static BigInteger getPublicModulus() {
        if (RSA.PUBLIC_MODULUS == null) {
            RSA.resetKeys();
        }
        return RSA.PUBLIC_MODULUS;
    }

    /**
     * Provides the exponent of the public key.
     *
     * @return {@link BigInteger} The exponent of the public key.
     */
    public static BigInteger getPublicExponent() {
        if (RSA.PUBLIC_EXPONENT == null) {
            RSA.resetKeys();
        }
        return RSA.PUBLIC_EXPONENT;
    }

    /**
     * Provides the modulus of the private key.
     *
     * @return {@link BigInteger} The modulus of the private key.
     */
    public static BigInteger getPrivateModulus() {
        if (RSA.PRIVATE_MODULUS == null) {
            RSA.resetKeys();
        }
        return RSA.PRIVATE_MODULUS;
    }

    /**
     * Provides the exponent of the private key.
     *
     * @return {@link BigInteger} The exponent of the private key.
     */
    public static BigInteger getPrivateExponent() {
        if (RSA.PRIVATE_EXPONENT == null) {
            RSA.resetKeys();
        }
        return RSA.PRIVATE_EXPONENT;
    }

    /**
     * Generates new values for the exponents and modulus of the keys.
     *
     * @param keySize An integer telling the size of the new keys in bytes. 2048
     * is recommended for users of this library.
     */
    public static void resetKeys(int keySize) {
        KeyPairGenerator kpg = null;
        try {
            kpg = KeyPairGenerator.getInstance("RSA");
        } catch (NoSuchAlgorithmException ex) {
            System.err.println("RSA algorithm couldn't be found in the system.");
        }
        kpg.initialize(keySize);
        KeyPair kp = kpg.genKeyPair();

        KeyFactory fact = null;
        try {
            fact = KeyFactory.getInstance("RSA");
        } catch (NoSuchAlgorithmException ex) {
            System.err.println("RSA algorithm couldn't be found in the system.");
        }
        RSAPublicKeySpec pub = null;
        try {
            pub = fact.getKeySpec(kp.getPublic(),
                    RSAPublicKeySpec.class);
        } catch (InvalidKeySpecException ex) {
            System.err.println(RSAPublicKeySpec.class + " was not accepted as a valid key specification.");
        }
        RSAPrivateKeySpec priv = null;
        try {
            priv = fact.getKeySpec(kp.getPrivate(),
                    RSAPrivateKeySpec.class);
        } catch (InvalidKeySpecException ex) {
            System.err.println(RSAPrivateKeySpec.class + " was not accepted as a valid key specification.");
        }

        RSA.PUBLIC_EXPONENT = pub.getPublicExponent();
        RSA.PUBLIC_MODULUS = pub.getModulus();
        RSA.PRIVATE_EXPONENT = priv.getPrivateExponent();
        RSA.PRIVATE_MODULUS = priv.getModulus();
    }

    /**
     * Generates new values for the exponents and modulus of the keys. The size
     * of the new keys will be 2048 bytes.
     */
    public static void resetKeys() {
        RSA.resetKeys(RSA.DEFAULT_KEY_SIZE);
    }

    /**
     * Encrypts a set of bytes by using the current class keys.
     *
     * @param bytes An array with the bytes to encrypt.
     * @return An array with the encrypted bytes.
     */
    public static byte[] encrypt(byte[] bytes) {
        return RSA.encrypt(bytes, RSA.PUBLIC_MODULUS, RSA.PUBLIC_EXPONENT);
    }

    /**
     * Decrypts a set of bytes by using the current class keys.
     *
     * @param bytes An array with the bytes to decrypt.
     * @return An array with the decrypted bytes.
     */
    public static byte[] decrypt(byte[] bytes) {
        return RSA.decrypt(bytes, RSA.PRIVATE_MODULUS, RSA.PRIVATE_EXPONENT);
    }

    /**
     * Encrypts a set of bytes by using custom modulus and exponent values
     * (therefore, custom keys).
     *
     * @param bytes An array with the bytes to encrypt.
     * @return An array with the encrypted bytes.
     */
    public static byte[] encrypt(byte[] bytes, BigInteger publicModulus, BigInteger publicExponent) {
        Cipher cipher = null;
        KeyFactory factory = null;
        try {
            cipher = Cipher.getInstance("RSA");
            factory = KeyFactory.getInstance("RSA");
        } catch (NoSuchAlgorithmException | NoSuchPaddingException ex) {
            if (ex instanceof NoSuchAlgorithmException) {
                System.err.println("RSA algorithm could not be found - " + ex.getMessage());
            } else if (ex instanceof NoSuchPaddingException) {
                System.err.println("Invalid padding scheme - " + ex.getMessage());
            }
        }
        try {
            cipher.init(Cipher.ENCRYPT_MODE, factory.generatePrivate(new RSAPrivateKeySpec(publicModulus, publicExponent)));
        } catch (InvalidKeyException | InvalidKeySpecException ex) {
            System.err.println(ex.getClass() + " when attempting to perform an encryption - " + ex.getMessage());
        }

        byte[] encrypted = null;
        try {
            encrypted = cipher.doFinal(bytes);
        } catch (IllegalBlockSizeException | BadPaddingException ex) {
            if (ex instanceof IllegalBlockSizeException) {
                System.err.println("Unable to process provided input data - " + ex.getMessage());
            } else if (ex instanceof BadPaddingException) {
                System.err.println("(Un)padding requested without corresponding padding bytes " + ex.getMessage());
            }
        }
        return encrypted;
    }

    /**
     * Decrypts a set of bytes by using custom modulus and exponent values
     * (therefore, custom keys).
     *
     * @param bytes An array with the bytes to decrypt.
     * @return An array with the decrypted bytes.
     */
    public static byte[] decrypt(byte[] bytes, BigInteger privateModulus, BigInteger privateExponent) {
        Cipher cipher = null;
        KeyFactory factory = null;
        try {
            cipher = Cipher.getInstance("RSA");
            factory = KeyFactory.getInstance("RSA");
        } catch (NoSuchAlgorithmException | NoSuchPaddingException ex) {
            if (ex instanceof NoSuchAlgorithmException) {
                System.err.println("RSA algorithm could not be found - " + ex.getMessage());
            } else if (ex instanceof NoSuchPaddingException) {
                System.err.println("Invalid padding scheme - " + ex.getMessage());
            }
        }
        try {
            cipher.init(Cipher.DECRYPT_MODE, factory.generatePrivate(new RSAPrivateKeySpec(privateModulus, privateExponent)));
        } catch (InvalidKeyException | InvalidKeySpecException ex) {
            System.err.println(ex.getClass() + " when attempting to perform a decryption - " + ex.getMessage());
        }

        byte[] decrypted = null;
        try {
            decrypted = cipher.doFinal(bytes);
        } catch (IllegalBlockSizeException | BadPaddingException ex) {
            if (ex instanceof IllegalBlockSizeException) {
                System.err.println("Unable to process provided input data - " + ex.getMessage());
            } else if (ex instanceof BadPaddingException) {
                System.err.println("(Un)padding requested without corresponding padding bytes " + ex.getMessage());
            }
        }
        return decrypted;
    }
}
